/*
 * Copyright (c) Huawei Technologies Co., Ltd. 2023-2023. All rights reserved.
 * Description:HCOM 基础函数
 * Author:luweijun
 * Create:2023-3-24
 */

#include "op_hcom_comm.h"
#include "hcom_pub.h"
#include "graph/ge_local_context.h"
#include "hcom_ops_kernel_info_store.h"
#include "sal.h"
#include "framework/common/ge_types.h" // ge对外options
#include "graph/utils/attr_utils.h"
#include "graph/utils/tensor_utils.h"
#include "graph/debug/ge_attr_define.h"
#include "graph/utils/node_utils.h"
#include "hcom_ops_kernel_info_store.h"
#include "hccl/hcom.h"
#include <nlohmann/json.hpp>
#include "externalinput.h"
namespace hccl {

HcclResult HcomInitialize()
{
    bool isInited = false;
    GetHcclCommInitStatus(isInited);
    if (isInited) {
        return HCCL_SUCCESS;
    }

    std::string rankTableAddrStr;
    std::string rankTableLenStr;
    ge::graphStatus geRet = ge::GetThreadLocalContext().GetOption(ge::OPTION_EXEC_RANK_TABLE_ADDR, rankTableAddrStr);
    CHK_PRT_RET((geRet != ge::GRAPH_SUCCESS),
        HCCL_WARNING("[HcomOpsKernelInfoStore][InitHcom]OPTION_EXEC_RANK_TABLE is not found."), HCCL_SUCCESS);
    geRet = ge::GetThreadLocalContext().GetOption(ge::OPTION_EXEC_RANK_TABLE_LEN, rankTableLenStr);
    CHK_PRT_RET((geRet != ge::GRAPH_SUCCESS),
        HCCL_ERROR("[HcomOpsKernelInfoStore][InitHcom]Failed to get OPTION_EXEC_RANK_TABLE_LEN."), HCCL_E_INTERNAL);

    u64 rankTableAddr;
    u32 rankTableLen;
    CHK_RET(SalStrToULonglong(rankTableAddrStr, HCCL_BASE_DECIMAL, rankTableAddr));
    CHK_RET(SalStrToULong(rankTableLenStr, HCCL_BASE_DECIMAL, rankTableLen));
    CHK_PRT_RET((rankTableLen > RANK_TABLE_MAX_LEN) || (rankTableLen == 0),
        HCCL_ERROR("[HcomOpsKernelInfoStore][InitHcom]errNo[0x%016llx] rankTable file is invalid, len is %u", \
        HCOM_ERROR_CODE(HCCL_E_PARA), rankTableLen), HCCL_E_PARA);
    std::string rankTable(reinterpret_cast<char *>(rankTableAddr), rankTableLen);

    std::string rankId;
    geRet = ge::GetThreadLocalContext().GetOption(ge::OPTION_EXEC_RANK_ID, rankId);
    CHK_PRT_RET((geRet != ge::GRAPH_SUCCESS),
        HCCL_ERROR("[HcomOpsKernelInfoStore][InitHcom]Failed to get OPTION_EXEC_RANK_ID."), HCCL_E_INTERNAL);

    // 该接口通过key值，获取ge session级别option的value
    std::string roleTableAddrStr;
    std::string roleTableLenStr;

    // 获取roleTable内存文件
    u32 roleTableLen = 0;
    geRet = ge::GetThreadLocalContext().GetOption(ge::OPTION_EXEC_ROLE_TABLE_LEN, roleTableLenStr);
    if (geRet == ge::GRAPH_SUCCESS) {
        CHK_RET(SalStrToULong(roleTableLenStr, HCCL_BASE_DECIMAL, roleTableLen));
    } else {
        HCCL_INFO("[HcomOpsKernelInfoStore][InitHcom]Option OPTION_EXEC_ROLE_TABLE_LEN was not found.");
    }
    if (roleTableLen == 0) {
        HCCL_INFO("[HcomOpsKernelInfoStore][InitHcom]OPTION_EXEC_ROLE_TABLE_LEN is zero.");
        CHK_RET(HcomInitByString(rankTable.c_str(), rankId.c_str()));
        return HCCL_SUCCESS;
    } else {
        CHK_RET(HcomInitByString(rankTable.c_str(), rankId.c_str(), WorkMode::HCCL_MODE_AI_CPU));
    }

    geRet = ge::GetThreadLocalContext().GetOption(ge::OPTION_EXEC_ROLE_TABLE_ADDR, roleTableAddrStr);
    CHK_PRT_RET((geRet != ge::GRAPH_SUCCESS),
        HCCL_ERROR("[HcomOpsKernelInfoStore][InitHcom]Failed to get OPTION_EXEC_ROLE_TABLE_ADDR."), HCCL_E_INTERNAL);

    u64 roleTableAddr;
    CHK_RET(SalStrToULonglong(roleTableAddrStr, HCCL_BASE_DECIMAL, roleTableAddr));
    CHK_PRT_RET((roleTableLen == 0),
        HCCL_INFO("[HcomOpsKernelInfoStore][InitHcom]roleTable file is invalid, len is 0."), HCCL_SUCCESS);
    CHK_PRT_RET((roleTableLen > RANK_TABLE_MAX_LEN),
        HCCL_ERROR("[HcomOpsKernelInfoStore][InitHcom]errNo[0x%016llx] roleTable file is invalid, len is %u", \
        HCOM_ERROR_CODE(HCCL_E_PARA), roleTableLen), HCCL_E_PARA);

    std::string reloTable(reinterpret_cast<char *>(roleTableAddr), roleTableLen);

    s32 rank = 0;
    HcclComm pComm = nullptr;
    CHK_RET(SalStrToInt(rankId, HCCL_BASE_DECIMAL, rank)); // 校验是否为有效值

    // 从Option中获取每条流中的表数量
    CHK_RET(GetEsMaxRemoteOpNumPerStreamFromGe());
    CHK_RET(HcclInitCsComm(rankTable.c_str(), rank, reloTable.c_str(), nullptr, &pComm));

    return HCCL_SUCCESS;
}

HcclResult GetEsMaxRemoteOpNumPerStreamFromGe()
{
    std::string esStreamNumStr{};
    ge::graphStatus geRet = ge::GetThreadLocalContext().GetOption(ge::OPTION_EXECUTE_TIMES,
        esStreamNumStr);
    CHK_PRT_RET((geRet != ge::GRAPH_SUCCESS),
        HCCL_ERROR("[HcomOpsKernelInfoStore][InitHcom]Failed to get OPTION_EXECUTE_TIMES."),
        HCCL_E_PARA);

    HCCL_INFO("esStreamNumStr: %s", esStreamNumStr.c_str());
    u32 esStreamNum;
    CHK_RET(SalStrToULong(esStreamNumStr, HCCL_BASE_DECIMAL, esStreamNum));
    HCCL_INFO("esStreamNum: %u", esStreamNum);
    SetEsStreamNum(esStreamNum);

    return HCCL_SUCCESS;
}

bool IsOpESScene()
{
    // 获取roleTable内存文件
    u32 roleTableLen = 0;
    std::string roleTableLenStr;

    ge::graphStatus geRet = ge::GetThreadLocalContext().GetOption(ge::OPTION_EXEC_ROLE_TABLE_LEN, roleTableLenStr);
    if (geRet == ge::GRAPH_SUCCESS) {
        CHK_RET(SalStrToULong(roleTableLenStr, HCCL_BASE_DECIMAL, roleTableLen));
    } else {
        HCCL_INFO("[HcomOpsKernelInfoStore][isOpESScene]Option OPTION_EXEC_ROLE_TABLE_LEN was not found.");
    }
    if (roleTableLen != 0) {
        return true;
    }
    return false;
}
HcclResult InitGroup()
{
    std::string groupConf;
    nlohmann::json groupListConf;
    ge::graphStatus geRet = ge::GetThreadLocalContext().GetOption(ge::OPTION_EXEC_HCOM_GROUPLIST, groupConf);
    CHK_PRT_RET((geRet != ge::GRAPH_SUCCESS),
        HCCL_WARNING("[HcomOpsKernelInfoStore][InitHcom]OPTION_EXEC_HCOM_GROUPLIST is not found."), HCCL_SUCCESS);
    HCCL_DEBUG("groupList:%s", groupConf.c_str());
    CHK_RET(SalParseInformation(groupListConf, groupConf));
    std::vector<nlohmann::json> groupList = groupListConf.get<std::vector<nlohmann::json>>();
    for (auto &groupInfo : groupList) {
        HCCL_DEBUG("groupInfo:%s", groupInfo.dump().c_str());
        std::vector<u32> ranks = groupInfo["group_rank_list"];
        std::string groupName = groupInfo["group_name"];
        HCCL_DEBUG("groupName:%s", groupName.c_str());

        u32 curRank = 0;
        CHK_RET(GetHcclRankId(HCCL_WORLD_GROUP, &curRank));
        if (!HcomFindGroup(groupName.c_str()) && find(ranks.begin(), ranks.end(), curRank) != ranks.end()) {
            if (!strncmp(groupName.c_str(), HCCL_WORLD_GROUP, sizeof(HCCL_WORLD_GROUP))) {
                HCCL_WARNING("[HcomOpsKernelInfoStore][InitHcom]cur groupname is HCCL_WORLD_GROUP.");
                continue;
            }
            CHK_RET(HcomCreateGroup(groupName.c_str(), ranks.size(), ranks.data()));
        }
    }
    return HCCL_SUCCESS;
}

HcclResult GetOpDescIntAttr(const ge::OpDesc &op, const string &attr, s32 &output)
{
    if (ge::AttrUtils::HasAttr(op, attr)) {
        ge::AttrUtils::GetInt(op, attr, output);
    } else {
        return HCCL_E_NOT_FOUND;
    }
    return HCCL_SUCCESS;
}

HcclResult GetOpDescStrAttr(const ge::OpDesc &op, const string &attr, string &output)
{
    if (ge::AttrUtils::HasAttr(op, attr)) {
        ge::AttrUtils::GetStr(op, attr, output);
    } else {
        return HCCL_E_NOT_FOUND;
    }
    return HCCL_SUCCESS;
}

HcclResult HcclGetRemoteOperationParams(const ge::OpDesc &opdesc, HcomRemoteOperationParams *opParams)
{
    CHK_PTR_NULL(opParams);
    /* get hcclType */
    opParams->hcclType = opdesc.GetType();
    /* get group */
    if (GetOpDescStrAttr(opdesc, "group", opParams->group) != HCCL_SUCCESS) {
        opParams->group = HCCL_WORLD_GROUP;
    }
    /* get "count" */
    if (GetOpDescIntAttr(opdesc, "max_num", opParams->count) != HCCL_SUCCESS) {
        HCCL_ERROR("[Get][RemoteOperationParams]get count from opdesc fail");
        return HCCL_E_INTERNAL;
    }
    /* get "valueDim" */
    if (GetOpDescIntAttr(opdesc, "embedding_dim", opParams->valueDim) != HCCL_SUCCESS) {
        HCCL_ERROR("[Get][RemoteOperationParams]get valueDim from opdesc fail");
        return HCCL_E_INTERNAL;
    }
    /* get "tag" */
    if (GetOpDescIntAttr(opdesc, "tag", opParams->tag) != HCCL_SUCCESS) {
        HCCL_ERROR("[Get][RemoteOperationParams]get tag from opdesc fail");
        return HCCL_E_INTERNAL;
    }
    /* get "insertOption" */
    s32 flags = 0;
    s32 insertFlag = 0;
    if ((opParams->hcclType == HCCL_KERNEL_OP_TYPE_COLL_REMOTE_LOOKUP ||
        opParams->hcclType == HCCL_KERNEL_OP_TYPE_COLL_REMOTE_LOOKUP_UNIQUED_PAIRED ||
        opParams->hcclType == HCCL_KERNEL_OP_TYPE_COLL_REMOTE_LOOKUP_PAIRED) &&
        (GetOpDescIntAttr(opdesc, "insert_option", insertFlag) != HCCL_SUCCESS) &&
        (GetOpDescIntAttr(opdesc, "flags", flags) != HCCL_SUCCESS)) {
        HCCL_ERROR("[Get][RemoteOperationParams]get insert_option or flags from opdesc fail");
        return HCCL_E_INTERNAL;
    }

    opParams->flag = static_cast<u64>(insertFlag);
    opParams->flags = flags;
    /* get "keyType" and "valueType" */
    CHK_RET(hccl::HcomOpUtils::TransformDataType(opdesc.GetInputDesc(1).GetDataType(), opParams->keyType));
    CHK_RET(hccl::HcomOpUtils::TransformDataType(opdesc.GetOutputDesc(0).GetDataType(), opParams->valueType));

    return HCCL_SUCCESS;
}

HcclResult GetHcclCommInitStatus(bool& initiated)
{
    CHK_RET(HcomGetInitStatus(initiated));

    if (!initiated) {
        CHK_RET(HcomRpcGetInitStatus(initiated));
    }
    return HCCL_SUCCESS;
}

HcclResult GetDeviceType(const char *group, DevType &deviceType)
{
    if (HcomIsRpcComm(group)) {
        HcomRpcGetDeviceType(group, deviceType);
    } else {
        std::shared_ptr<hccl::hcclComm> hcclComm;
        CHK_RET(HcomGetCommByGroup(group, hcclComm));
        CHK_RET(hcclComm->GetDevType(deviceType));
    }

    return HCCL_SUCCESS;
}

HcclResult GetRankSize(const char *group, u32 *rankSize)
{
    if (HcomIsRpcComm(group)) {
        CHK_RET(HcomRpcGetRankSize(group, rankSize));
    } else {
        CHK_RET(HcomGetRankSize(group, rankSize));
    }
    return HCCL_SUCCESS;
}

HcclResult GetHccsLinkNum(const char *group, u32 &numHccsLink)
{
    if (HcomIsRpcComm(group)) {
        CHK_RET(HcomRpcGetHccsLinkNum(group, numHccsLink));
    } else {
        CHK_RET(HcomGetHccsLinkNum(group, numHccsLink));
    }
    return HCCL_SUCCESS;
}

HcclResult GetHcclRankId(const char *group, u32 *rankId)
{
    if (HcomIsRpcComm(group)) {
        CHK_RET(HcomRpcGetRankId(group, rankId));
    } else {
        CHK_RET(HcomGetRankId(group, rankId));
    }
    return HCCL_SUCCESS;
}

HcclResult UnloadHcclTask(const std::string group, const char *tag)
{
    if (HcomIsRpcComm(group.c_str())) {
        CHK_RET(HcomRpcUnloadTask(group, tag));
    } else {
        CHK_RET(HcomUnloadTask(group, tag));
    }
    return HCCL_SUCCESS;
}

HcclResult SetWorkspaceResource(const std::string &tag, const char *group, std::vector<rtStream_t> stream,
    void *memPtr, u64 maxSize)
{
    if (HcomIsRpcComm(group)) {
        CHK_RET(HcomRpcSetWorkspaceResource(tag, group, stream, memPtr, maxSize));
    } else {
        CHK_RET(HcomSetWorkspaceResource(tag, group, stream, memPtr, maxSize));
    }
    return HCCL_SUCCESS;
}

HcclResult GetTopoDesc(const char *group, HcclTopoDescs *topoDescs, uint32_t topoSize)
{
    if (HcomIsRpcComm(group)) {
        CHK_RET(HcomRpcGetTopoDesc(group, topoDescs, topoSize));
    } else {
        CHK_RET(HcomGetTopoDesc(group, topoDescs, topoSize));
    }
    return HCCL_SUCCESS;
}
}